package com.qzdatasoft.ratelimit.adapter;

import com.qzdatasoft.ratelimit.bean.ResourceKey;
import com.qzdatasoft.ratelimit.bean.RateLimitRule;
import com.qzdatasoft.ratelimit.po.RateLimitRulePo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 限流规则缓存类
 */
public class RateLimitRuleCacheManager implements InitializingBean {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitRuleCacheManager.class);

    private static final AtomicReference<Map<String, RateLimitRule>> ruleCache =
            new AtomicReference<>(new ConcurrentHashMap<>());
    private final RateLimitRuleQueryProvider queryProvider;
    private final ClientInfoExtractor clientInfoExtractor;

    public RateLimitRuleCacheManager(
            RateLimitRuleQueryProvider queryProvider, ClientInfoExtractor clientInfoExtractor
    ) {
        this.queryProvider = queryProvider;
        this.clientInfoExtractor = clientInfoExtractor;
    }

    public void initializeCache() {
        Map<String, RateLimitRule> initialCache = new ConcurrentHashMap<>();

        for (RateLimitRulePo rule : queryProvider.loadEnabledRules()) {
            if (this.isOwner(rule.getUrlPattern())) {
                try(ResourceKey resourceKey = ResourceKey.parseKey(rule.getUrlPattern())) {
                        initialCache.put(rule.getUrlPattern(), RateLimitRule.builder()
                                    .urlPattern(resourceKey.getUrlPattern())
                                    .dimension(resourceKey.getDimension())
                                    .resourceKey(rule.getUrlPattern())
                                    .user(resourceKey.getUser())
                                    .ip(resourceKey.getIp())
                                    .burstCapacity(rule.getBurstCapacity())
                                    .compensationRate(rule.getCompensationRate())
                                    .timeoutMillis(rule.getTimeoutMillis()) // 注意超时时间单位为毫秒
                                    .ratePerSecond(rule.getRatePerSecond())
                                    //.nodelay(false)
                                    .build()
                    );
                    logger.debug("Added/Updated rate limit rule: {}", rule.getId());
                } catch (Exception e){
                    logger.error("Error parsing resource key: {}", rule.getUrlPattern(), e);
            }
            }
        }

        ruleCache.set(initialCache);
        logger.info("Initialized rate limit rule cache with {} rules", initialCache.size());
    }

    public RateLimitRule getRule(String rulPattern) {
        return ruleCache.get().get(rulPattern);
    }

    public Map<String, RateLimitRule> getAllRules() {
        return new ConcurrentHashMap<>(ruleCache.get());
    }

    public void putRule(String rulPattern, RateLimitRule rule) {
        if (rule != null && rule != null & this.isOwner(rulPattern)) {
            ruleCache.get().put(rulPattern, rule);
            logger.info("Added/Updated rate limit rule from cache: {}", rulPattern);
        }
    }

    public RateLimitRule removeRule(String rulPattern) {
        if (rulPattern != null) {
            RateLimitRule removed = ruleCache.get().remove(rulPattern);
            if (removed != null) {
                logger.info("Removed rate limit rule from cache: {}", rulPattern);
            } else {
                logger.warn("Rate limit rule not found for removal from cache: {}", rulPattern);
            }
            return removed;
        }
        return null;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        initializeCache();
    }

    private boolean isOwner(String rulPattern) {
        return clientInfoExtractor.isOwner(rulPattern);
    }

}